#coding=utf-8
import jieba.posseg as pseg
from DBConnection import DBConnection


class Snippet(object):

    def __init__(self):
        """链接数据库"""
        self.connection = DBConnection('news', 'wangyi_news')
        self.connection.connection()

    def create_snippet(self, query, news_id, window=100, query_list=None):
        """
        snippet生成
        :param query: 用户输入的查询 (str)
        :param news_id: 该查询对应的文档id号(list)
        :param window: snippet 窗口大小, 默认为100
        :param query_list: 用户输入的查询转换成list（可省）
        :return: 生成的snippet [str, str, str...]
        """
        # 如果包含完整的query
        flag = False
        time_list = self.connection.get_news_time(news_id)
        content_list = self.connection.get_news_content(news_id)

        news_snippet = []
        for k in range(len(content_list)):
            content = content_list[k]
            index = 1
            for i in range(1, len(content)):
                if query in content[i]:
                    index = i
                    flag = True
                    break

            # 如果不包含完整的query
            if not flag:
                if query_list is None:  # 如果query_list为空，则自己分词
                    query_list = []
                    stopwords = ['uj', 'ul', 'x', 'eng']
                    words = pseg.cut(query)
                    for word, tag in words:
                        if tag not in stopwords:
                            query_list.append(word)

                # 获取包含最多关键词的句子作为snippet的首句
                max_count = 0
                for i in range(1, len(content)):
                    count = 0
                    for query_word in query_list:
                        if query_word in content[i]:
                            count += 1
                    if count > max_count:
                        index = i

            # snippet 生成
            snippet = content[index]
            while len(snippet) < window and index + 1 < len(content) - 1:
                index += 1
                snippet += content[index]
            news_snippet.append(time_list[k] + ' - ' + snippet[0:window] + '...')
        return news_snippet

    def create_preview(self, news_id):
        """
        生成的预览
        :param news_id: 该查询对应的文档id号(list)
        :return: 生成的预览 [[title, content], [title, content], ... ]
        """
        title_list = self.connection.get_news_title(news_id)
        content_list = self.connection.get_news_content(news_id)
        news_preview = []
        for i in range(len(content_list)):
            news_preview.append([title_list[i], ' '.join(content_list[i][1:])])
        return news_preview


if __name__ == '__main__':
    snippet = Snippet()

    # 生成snippet
    # news_id = ['D55N97190001875N', 'D54RFIQV0001875N']
    news_id = [u'D54T85S30001875N']

    news_snippet = snippet.create_snippet('信息检索'.decode("utf-8"), news_id)  # 生成的snippet [str, str, str...]
    print(news_snippet)

    # 生成预览
    preview = snippet.create_preview(news_id)  # 生成的预览 [[title, content], [title, content], ... ]
    print(preview)


